import os
import cv2
import numpy as np
try:
    from .cvio import cvio
except:
    from cvio import cvio


def rectify_fridge_images(src, dst=''):
    mtx = np.array([[636.7940420735532, 0.0, 637.0192231645837], [0.0, 642.2078175888157, 312.67232337362935], [0.0, 0.0, 1.0]])
    dist = np.array([[-0.36721049616656565, 0.11686644036526267, 0.0013678611194601928, 0.0013781249474976913, -0.0159662883808598]])
    images = cvio.load_image_list(src,recursive=True,silent=False)
    n = len(images)
    if dst == '' or dst is None:
        dst = src
    if not os.path.exists(dst):
        os.makedirs(dst)
    for i,impath in enumerate(images,1):
        print('[%d/%d] %s' % (i, n, os.path.basename(impath)))
        img = np.array(cvio.load_img(impath, mode='pil'))
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
        res = cv2.undistort(img, mtx, dist, None, None)
        res = cv2.rotate(res, cv2.ROTATE_90_COUNTERCLOCKWISE)
        imsave = impath if dst == src else os.path.join(dst,os.path.basename(impath))
        cvio.write_img(res[:,:,::-1], imsave)


def adjust_gamma_images(src, minvalue=0, maxvalue=211, gamma=1.9):
    imglist = cvio.load_image_list(src,recursive=True,silent=False)
    n = len(imglist)

    if n <= 0: 
        print('无发现图片数据！')
        return

    # img adjust para
    inBlack = np.array([minvalue, minvalue, minvalue], dtype=np.float32)
    inWhite = np.array([maxvalue, maxvalue, maxvalue], dtype=np.float32)
    in_between = inWhite - inBlack
    inGamma = np.array([gamma, gamma, gamma], dtype=np.float32)
    inv_inGamma = 1 / inGamma
    outBlack = np.array([0, 0, 0], dtype=np.float32)
    outWhite = np.array([255, 255, 255], dtype=np.float32)

    def adjust_gamma_impl(img):
        img = np.clip((img) / in_between, 0, 255)
        img = (img ** inv_inGamma) * outWhite
        img = np.clip(img, 0, 255).astype(np.uint8)
        return img

    for i,imp in enumerate(imglist,1):
        print('[%d/%d] %s' % (i,n,os.path.basename(imp)))
        img = np.array(cvio.load_img(imp, mode='pil'), dtype=np.uint8)
        img = adjust_gamma_impl(img)
        cvio.write_img(img, imp, mode='pil')

